#include "slice2cube.h"

#include <cstring>
#include <thread>
#include <fstream>

#include <sys/stat.h>

#include "public/utils.h"

void createDirectory(const std::string& dir) {
	struct stat buf;
	auto dirc=dir.c_str();
	int r=stat(dirc, &buf);
	if(r!=-1 && S_ISDIR(buf.st_mode))
		return;
#ifdef BUILD_MINGW
	r=mkdir(dirc);
#else
	r=mkdir(dirc, 0777);
#endif
	if(r!=-1)
		return;
	throwError("Cannot create dir '", dir, "': ", strerror(errno), ".\n");
}

template <typename C>
inline std::ostream& operator<<(std::ostream& fs, const std::chrono::time_point<C>& tp) {
	auto tt=std::chrono::system_clock::to_time_t(tp);
	char buf[64];
	auto r=strftime(buf, 64, "%Y-%m-%d %H:%M:%S", localtime(&tt));
	fs.write(buf, r);
	return fs;
}

inline std::ostream& operator<<(std::ostream& fs, const std::chrono::milliseconds& dur) {
	auto v=dur.count();
	fs.width(2);
	fs.fill('0');
	fs<<v/1000/3600<<':';
	v=v%(1000*3600);
	fs.width(2);
	fs.fill('0');
	fs<<v/1000/60<<':';
	v=v%(1000*60);
	fs.width(2);
	fs.fill('0');
	fs<<v/1000<<'.';
	v=v%1000;
	fs.width(3);
	fs.fill('0');
	fs<<v;
	return fs;
}

template <typename T>
Slice2Cube<T>::Slice2Cube(const std::vector<std::string>& _inputfiles, std::vector<std::pair<const std::string*, int>>&& _inputpairs, double _xres, double _yres, double _zres, uint64_t _width, uint64_t _height, uint64_t _depth, int _channels, int _bps, uint64_t _xsize, uint64_t _ysize, uint64_t _zsize, uint32_t _dsx, uint32_t _dsy, uint32_t _dsz, const std::string& _outdir, int64_t _cachesize, const uint8_t* _xfunc, int _ncpu, bool _resume, bool _ds_only, bool _ds_max):
	inputfiles(_inputfiles), xres{_xres}, yres{_yres}, zres{_zres},
	width{_width}, height{_height}, depth{_depth},
	channels{_channels}, bps{_bps},
	xsize{_xsize}, ysize{_ysize}, zsize{_zsize},
	dsw{0}, dsh{0}, dsd{0},
	dsx{_dsx}, dsy{_dsy}, dsz{_dsz},
	outdir(_outdir), cachesize{_cachesize}, xfunc{_xfunc},
	ncpu{_ncpu}, resume{_resume},
	ds_only{_ds_only}, ds_max{_ds_max},
	all_checked{_inputpairs.size()>0}, inputpairs(std::move(_inputpairs)),
	tiffidx{0}, tiffd{nullptr},
	imagebufs{}, tempfs(channels), nrrdfs(channels), tempfn{},
	start_idx_ds{0},
	bufds{}, bufdssum{}, writers{},
	error{false}, finished{false},
	mtx{}, cv0{}, cv1{},
	readQue{}, writeQue{}, writeCubeQue{}, writeDsQue{},
	readQueSize{0}, writeQueSize{0}, writeCubeQueSize{0}, writeDsQueSize{0}
{
	if(dsx==1 && dsy==1 && dsz==1)
		ds_only=true;
	if(xsize<=0) {
		uint64_t xn=1, yn=1, zn=1;
		xsize=((width+xn-1)/xn+7)/8*8;
		ysize=((height+yn-1)/yn+7)/8*8;
		zsize=((depth+zn-1)/zn+7)/8*8;
		while(xsize*ysize*zsize>16*1024*1024) {
			int maxis=0;
			double mlen=xsize*xres;
			if(ysize*yres>mlen) {
				mlen=ysize*yres;
				maxis=1;
			}
			if(zsize*zres>mlen) {
				mlen=zsize*zres;
				maxis=2;
			}
			if(maxis==0) {
				xn++;
			} else if(maxis==1) {
				yn++;
			} else {
				zn++;
			}
			xsize=((width+xn-1)/xn+7)/8*8;
			ysize=((height+yn-1)/yn+7)/8*8;
			zsize=((depth+zn-1)/zn+7)/8*8;
		}
		if(zn==1)
			zsize=depth;
	}
	if(dsx<=0) {
		dsx=1;
		dsy=1;
		dsz=1;
		dsw=(width+dsx-1)/dsx;
		dsh=(height+dsy-1)/dsy;
		dsd=(depth+dsz-1)/dsz;
		while(dsw*dsh*dsd>320*1024*1024) {
			int maxis=0;
			auto mres=dsx*xres;
			if(dsy*yres<mres) {
				mres=dsy*yres;
				maxis=1;
			}
			if(dsz*zres<mres) {
				mres=dsz*zres;
				maxis=2;
			}
			if(maxis==0) {
				dsx++;
			} else if(maxis==1) {
				dsy++;
			} else {
				dsz++;
			}
			dsw=(width+dsx-1)/dsx;
			dsh=(height+dsy-1)/dsy;
			dsd=(depth+dsz-1)/dsz;
		}
	} else {
		dsw=(width+dsx-1)/dsx;
		dsh=(height+dsy-1)/dsy;
		dsd=(depth+dsz-1)/dsz;
	}
}
template Slice2Cube<uint8_t>::Slice2Cube(const std::vector<std::string>& _inputfiles, std::vector<std::pair<const std::string*, int>>&& _inputpairs, double _xres, double _yres, double _zres, uint64_t _width, uint64_t _height, uint64_t _depth, int _channels, int _bps, uint64_t _xsize, uint64_t _ysize, uint64_t _zsize, uint32_t _dsx, uint32_t _dsy, uint32_t _dsz, const std::string& _outdir, int64_t _cachesize, const uint8_t* _xfunc, int _ncpu, bool _resume, bool _ds_only, bool _ds_max);
template Slice2Cube<uint16_t>::Slice2Cube(const std::vector<std::string>& _inputfiles, std::vector<std::pair<const std::string*, int>>&& _inputpairs, double _xres, double _yres, double _zres, uint64_t _width, uint64_t _height, uint64_t _depth, int _channels, int _bps, uint64_t _xsize, uint64_t _ysize, uint64_t _zsize, uint32_t _dsx, uint32_t _dsy, uint32_t _dsz, const std::string& _outdir, int64_t _cachesize, const uint8_t* _xfunc, int _ncpu, bool _resume, bool _ds_only, bool _ds_max);

template <typename T>
void Slice2Cube<T>::printConversionInfo() {
	std::cerr<<"Input images: "<<bps<<" bit, "<<channels<<(channels>1?" channels, ":" channel, ")<<width<<"*"<<height<<"*"<<depth<<" voxels";
	std::cerr<<" ("<<width*xres<<"*"<<height*yres<<"*"<<depth*zres<<" µm³)\n";
	if(!ds_only) {
		std::cerr<<"Output cubes: "<<sizeof(T)*8<<" bit, "<<xsize<<"*"<<ysize<<"*"<<zsize<<" voxels";
		std::cerr<<" ("<<xsize*xres<<"*"<<ysize*yres<<"*"<<zsize*zres<<" µm³)\n";
	}
	std::cerr<<"Downsampled volume [";
	if(ds_max)
		std::cerr<<"MIP";
	else
		std::cerr<<"MEAN";
	std::cerr<<"]: "<<dsw<<"*"<<dsh<<"*"<<dsd<<" voxels";
	std::cerr<<" (factors "<<dsx<<":"<<dsy<<":"<<dsz<<")\n";
	auto cubesize=xsize*ysize*zsize*sizeof(T);
	if(cubesize*3*3*3>1024*1024*1024)
		std::cerr<<"WARN: Output cubes might be too big! ("<<cubesize*3*3*3/1024.0/1024/1024<<" GiB for 3*3*3 cubes)\n";
	if(cubesize<4*1024*1024)
		std::cerr<<"WARN: Output cubes might be too small! ("<<cubesize/1024.0/1024<<" MiB per cube)\n";
	auto dssize=dsw*dsh*dsd*sizeof(T);
	if(dssize>512*1024*1024)
		std::cerr<<"WARN: Downsampled volume might be too big! ("<<dssize/1024.0/1024<<" MiB)\n";
}
template void Slice2Cube<uint8_t>::printConversionInfo();
template void Slice2Cube<uint16_t>::printConversionInfo();

#if 0
{
	writeCatalogSimple();
	auto time_end=std::chrono::system_clock::now();
	time_delta=std::chrono::duration_cast<std::chrono::milliseconds>(time_end-time_start);
	std::cerr<<"["<<time_delta<<"] Complete ("<<time_end<<").\n";
}
#endif
template <typename T>
void Slice2Cube<T>::writeCatalogSimple() {
	std::ofstream fs;
	std::ostringstream oss;
	oss<<outdir<<"/catalog";
	auto fn=oss.str();
	fs.open(fn);
	if(!fs)
		throwError("Failed to open file: ", fn, ".");
	for(int i=0; i<channels; i++) {
		fs<<"[CH";
		fs.width(2);
		fs.fill('0');
		fs<<i<<"]\n";
		fs<<"location=ch";
		fs.width(2);
		fs.fill('0');
		fs<<i<<".nrrd\n";
		fs<<"origin=0 0 0\ndirection="<<xres<<" 0 0 0 "<<yres<<" 0 0 0 "<<zres<<"\n";
		if(!fs)
			throwError("Failed to write to file.");
	}
	fs.close();
	if(!fs)
		throwError("Failed to close file: ", fn, ".");
}
template <typename T>
void Slice2Cube<T>::startConvertToCubesWorker() {
	try {
		std::unique_lock<std::mutex> lck{mtx};
		while(!error && !finished) {
			if(!writeDsQue.empty()) {
				auto slices=&writeDsQue.front();
				if(slices->scheduled<slices->total) {
					auto jobi=slices->scheduled;
					slices->scheduled++;
					lck.unlock();
					if(slices->bufs.size()==0) {
						deflateDownsample(slices, jobi);
					} else {
						writeDownsample(slices, jobi);
					}
					lck.lock();
					slices->unfinished--;
					if(slices->unfinished==0)
						cv0.notify_one();
					continue;
				}
			}
			if(!writeCubeQue.empty()) {
				auto slices=&writeCubeQue.front();
				if(slices->scheduled<slices->total) {
					auto jobi=slices->scheduled;
					slices->scheduled++;
					if(!ds_only) {
						lck.unlock();
						writeCubes(slices, jobi);
						lck.lock();
					}
					slices->unfinished--;
					if(slices->unfinished==0)
						cv0.notify_one();
					continue;
				}
			}
			if(!readQue.empty()) {
				auto slices=readQue.begin();
				while(slices!=readQue.end() && slices->scheduled>=slices->total)
					++slices;
				if(slices!=readQue.end()) {
					auto jobi=slices->scheduled;
					slices->scheduled++;
					lck.unlock();
					readInputFiles(&(*slices), jobi);
					lck.lock();
					slices->unfinished--;
					if(slices->unfinished==0)
						cv0.notify_one();
					continue;
				}
			}
			cv1.wait(lck);
		}
	} catch(std::exception& e) {
		std::unique_lock<std::mutex> lck{mtx};
		printMessage("Error1: ", e.what());
		error=true;
		cv0.notify_one();
	}
}

template <typename T>
void Slice2Cube<T>::getResumePosition(uint64_t* pstart_idx) {
	uint64_t start_idx=depth;
	start_idx_ds=depth;

	if(resume) {
		for(int i=0; i<channels; i++) {
			struct stat buf;
			auto r=stat(tempfn[i].c_str(), &buf);
			if(r==-1 || !S_ISREG(buf.st_mode))
				throwError("Some temporary files missing. Cannot continue.\nConversion is either not started or already finished.\n");
			uint64_t z=buf.st_size/sizeof(T)/dsw/dsh;
			if(z*dsz<start_idx_ds)
				start_idx_ds=z*dsz;
		}
		if(!ds_only)
			start_idx=start_idx_ds/zsize*zsize;
		for(int i=0; i<channels; i++) {
			tempfs[i].open(tempfn[i], std::ofstream::out|std::ofstream::binary|std::ofstream::app);
			if(!tempfs[i])
				throwError("Failed to open file: ", tempfn[i], ".");
			tempfs[i].seekp(start_idx_ds/dsz*dsw*dsh*sizeof(T), std::ofstream::beg);
			if(!tempfs[i])
				throwError("Failed to seekp file: ", tempfn[i], ".");
		}
	} else {
		start_idx=start_idx_ds=0;
		for(int i=0; i<channels; i++) {
			tempfs[i].open(tempfn[i], std::ofstream::out|std::ofstream::binary|std::ofstream::trunc);
			if(!tempfs[i])
				throwError("Failed to open file: ", tempfn[i], ".");
		}
	}
	*pstart_idx=start_idx;
}

template <typename T>
void Slice2Cube<T>::startConvertToCubes() {
	tempfn.resize(channels);
	for(int i=0; i<channels; i++) {
		std::ostringstream oss;
		oss<<outdir<<"/ch";
		oss.width(2);
		oss.fill('0');
		oss<<i;
		auto dir=oss.str();
		oss<<"ds.temp";
		tempfn[i]=oss.str();
		if(!ds_only) {
			createDirectory(dir);
			for(uint64_t z=0; z<depth; z+=zsize) {
				std::ostringstream oss;
				oss<<dir<<"/z";
				oss.width(8);
				oss.fill('0');
				oss<<z;
				auto dirname=oss.str();
				createDirectory(dirname);
			}
		}
	}

	uint64_t start_idx;
	getResumePosition(&start_idx);
	if(start_idx!=0) {
		printMessage("Resuming from ", start_idx);
	}

	bufds.resize(dsw*dsh*channels);
	bufdssum.resize(dsw*dsh*channels);

	uint64_t zinserted=0;
	while(zinserted<start_idx) {
		if(!all_checked)
			getNextImage();
		zinserted++;
	}
	uint64_t zfinished=start_idx;

	for(int i=0; i<ncpu; i++) {
		imagebufs.push_back({});
		if(imagebufs.size()>=depth)
			break;
	}

	if(!ds_only)
		writers.resize(channels*((width+xsize-1)/xsize)*((height+ysize-1)/ysize));

	std::vector<std::thread> threads{};
	for(int i=0; i<ncpu; i++)
		threads.emplace_back(&Slice2Cube<T>::startConvertToCubesWorker, this);

	auto time_start=std::chrono::system_clock::now();
	auto time_delta=std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now()-time_start);
	std::cerr<<"["<<time_delta<<"] Begin converting slices ("<<time_start<<").\n";
	std::unique_lock<std::mutex> lck{mtx};

	for(size_t bi=0; bi<imagebufs.size(); bi++) {
		if(zinserted>=depth)
			break;
		imagebufs[bi].resize(width*height*channels);
		readQue.push_back({{imagebufs[bi].data()}, zinserted, ncpu, ncpu, ncpu});
		auto slices=&readQue.back();
		lck.unlock();
		if(!all_checked)
			getNextImage();
		zinserted++;
		lck.lock();
		slices->scheduled=0;
		readQueSize+=slices->bufs.size();
		cv1.notify_all();
	}

	while(true) {
		if(error) {
			cv1.notify_all();
			lck.unlock();
			printMessage("joining threads");
			for(auto& t: threads)
				t.join();
			throwError("Error in worker threads.");
		}
		if(!writeDsQue.empty()) {
			auto slices=&writeDsQue.front();
			if(slices->unfinished==0) {
				auto bufs=slices->bufs;
				writeDsQueSize-=bufs.size();
				writeDsQue.pop_front();
				time_delta=std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now()-time_start);
				printMessage("[", time_delta, "] Slices [", zfinished, ", ", zfinished+bufs.size(), ") finished.");
				zfinished+=bufs.size();
				if(zfinished>=depth)
					break;
				cv1.notify_all();
				for(size_t bi=0; bi<bufs.size(); bi++) {
					if(zinserted>=depth)
						break;
					readQue.push_back({{bufs[bi]}, zinserted, ncpu, ncpu, ncpu});
					auto slices1=&readQue.back();
					lck.unlock();
					if(!all_checked)
						getNextImage();
					zinserted++;
					lck.lock();
					slices1->scheduled=0;
					readQueSize+=slices1->bufs.size();
					cv1.notify_all();
				}
				continue;
			}
		}
		if(!readQue.empty()) {
			auto slices=&readQue.front();
			auto slicesi=readQue.begin();
			if(slices->unfinished==0 && (static_cast<size_t>((ncpu+3)/4-1)>=readQue.size() || (slicesi+(ncpu+3)/4-1)->unfinished==0)) {
				writeCubeQue.push_back({{}, slices->zidx, 0, 0, 0});
				auto slices1=&writeCubeQue.back();
				do {
					for(auto b: slices->bufs)
						slices1->bufs.push_back(b);
					readQueSize-=slices->bufs.size();
					readQue.pop_front();
					if(readQue.empty()) {
						slices=nullptr;
					} else {
						slices=&readQue.front();
					}
				} while(slices && slices->unfinished==0);
				slices1->scheduled=0;
				slices1->unfinished=slices1->total=channels*ncpu*ncpu;
				writeCubeQueSize+=slices1->bufs.size();
				cv1.notify_all();
				continue;
			}
		}
		if(!writeCubeQue.empty()) {
			auto slices=&writeCubeQue.front();
			if(slices->unfinished==0) {
				writeDsQue.push_back({{}, slices->zidx, 0, 0, 0});
				auto slices1=&writeDsQue.back();
				do {
					for(auto b: slices->bufs)
						slices1->bufs.push_back(b);
					writeCubeQueSize-=slices->bufs.size();
					writeCubeQue.pop_front();
					if(writeCubeQue.empty()) {
						slices=nullptr;
					} else {
						slices=&writeCubeQue.front();
					}
				} while(slices && slices->unfinished==0);
				slices1->scheduled=0;
				slices1->unfinished=slices1->total=channels;
				writeDsQueSize+=slices1->bufs.size();
				cv1.notify_all();
				continue;
			}
		}
		cv0.wait(lck);
	}
	{
		for(int i=0; i<channels; i++) {
			tempfs[i].close();
			if(!tempfs[i])
				throwError("Failed to close file.");
		}
		writeDsQue.push_back({{}, depth, channels, channels, channels});
		auto slices1=&writeDsQue.back();
		slices1->scheduled=0;
		writeDsQueSize+=1;
		cv1.notify_all();
		while(slices1->unfinished!=0)
			cv0.wait(lck);
	}
	finished=true;
	cv1.notify_all();
	lck.unlock();

	for(auto& t: threads)
		t.join();

	writeCatalogCubes();

	for(int i=0; i<channels; i++) {
		if(remove(tempfn[i].c_str())==-1)
			throwError("Failed to remove temporary file");
	}
	auto time_end=std::chrono::system_clock::now();
	time_delta=std::chrono::duration_cast<std::chrono::milliseconds>(time_end-time_start);
	std::cerr<<"["<<time_delta<<"] Complete ("<<time_end<<").\n";
}

template <typename T>
void Slice2Cube<T>::writeCatalogCubes() {
	std::ofstream fs;
	std::ostringstream oss;
	oss<<outdir<<"/catalog";
	auto fn=oss.str();
	fs.open(fn);
	if(!fs)
		throwError("Failed to open file: ", fn, ".");
	for(int i=0; i<channels; i++) {
		fs<<"[CH";
		fs.width(2);
		fs.fill('0');
		fs<<i<<"DS]\n";
		fs<<"location=ch";
		fs.width(2);
		fs.fill('0');
		fs<<i<<"ds.nrrd\n";
		fs<<"origin=0 0 0\ndirection="<<xres*dsx<<" 0 0 0 "<<yres*dsy<<" 0 0 0 "<<zres*dsz<<"\n";
		if(!ds_only) {
			fs<<"[CH";
			fs.width(2);
			fs.fill('0');
			fs<<i<<"]\n";
			fs<<"pattern=ch";
			fs.width(2);
			fs.fill('0');
			fs<<i;
			fs<<"/z<08Z>/y<08Y>.x<08X>.nrrd\n";
			fs<<"size="<<width<<" "<<height<<" "<<depth<<"\n";
			fs<<"cubesize="<<xsize<<" "<<ysize<<" "<<zsize<<"\n";
			fs<<"origin=0 0 0\ndirection="<<xres<<" 0 0 0 "<<yres<<" 0 0 0 "<<zres<<"\n";
		}
		if(!fs)
			throwError("Failed to write to file.");
	}
	fs.close();
	if(!fs)
		throwError("Failed to close file: ", fn, ".");
}

template void Slice2Cube<uint8_t>::startConvertToCubes();
template void Slice2Cube<uint16_t>::startConvertToCubes();
